#!/usr/bin/python
#
# Script for displaying run-time function call information in the form of
# nested function code.
#
# To install 'whackme' as an lldb command, at the (lldb) prompt enter:
#    command script import <pathto>/stackwhack.py
#
# See http://llvm.org/viewvc/llvm-project/lldb/trunk/include/lldb/API/
# for information about the Python lldb API.
#
# Miranda R. Callahan

# TODO: remove imports not used.
import lldb
import commands
import optparse
import shlex
import collections
import re

# Version!
print "version 3.1"

# Name of the file to which to print the unwound runtime stack.
stackfilename = 'stackfile'

def sw():
  """ Call stackwhack on the current thread:
  """
  stackwhack(lldb.frame.GetThread())

def getReturnType(type):
  """ Given an SBType, get only the return type.
      This function isn't included in our version of SBFunction. :-P

      Args:
        type: an SBType

      Returns:
        a string representing the return type only from |type|.
  """
  rettype = type.split()[0]
  # Start hacking in objc possibilities here. :-P
  if (rettype == "void"):
    rettype = "(void|IBAction)"
  return rettype

# Data structure so we can return the start and end lines of a function.
Funcbounds = collections.namedtuple('Funcbounds', ['start', 'end'],
                                    verbose=False)

def getFuncBounds(funname, funtype, filename, includesline):
  """ Takes information about a function and returns the starting and ending
      line where the function can be found.

      Args:
        funname: Gives the function name as a string.
        funtype: Gives the function return type as an SBType.
        filename: Gives the absolute file path and name in which the function
            can be found. 
        includesline: Gives a line inside the function.

      Returns: 
        A Funcbounds tuple giving the start and end lines of the function.
        Lines are zero-based (so the first line in the file is line 0). This
        makes it easier to index into the file stored as an array of lines.
  """
  funfile = open(filename, 'r')
  filelines = funfile.read().splitlines()
  # TODO: in the case of objective c, the type must be mangled to have parens,
  #       and - or +. OR use a regex?
  #       in the case of c++ (and others??), the name must be stripped of
  #       params. (funname.split("(")[0])
  print "**************\n"
  findfunPattern = funname  # Just a holder for now.
  print "initial funname: " + funname
 
  # Do more splitting if this is an objective C function:
  objc_funparts = funname.split("[")
  if (len(objc_funparts) > 1):  # TODO: can this ever be more than 2?
    newfunname = objc_funparts[1].split("]")[0].split(" ")[1]
    # now strip all but first part of method name. HACK!!! NOOO!!
    newfunname = newfunname.split(":")[0]
    findfunPattern = funname[0] + " \(" + getReturnType(funtype) + "\)" + newfunname
    funname = newfunname
    print "OBJECTIVE C FOUND:\n"
  else:
    funname = funname.split("(")[0]  # Remove params for non-objective-c.
    findfunPattern = getReturnType(funtype) + " " + funname
    print "C CODE FOUND:\n"

  print funname
  print findfunPattern

  print "**************\n"
  startfun = -1 
  ix = 0
  assert (len(filelines) >= 
          includesline), "ERROR in getFuncBounds::includesline."
  # Get the last line between 0 and includesline that starts with the function
  # type and name.
  while (ix <= includesline):
    if ((includesline - ix) < 100):
      print filelines[ix]
    if (re.match(findfunPattern, filelines[ix])):
      startfun = ix  
    ix = ix + 1
  # If we haven't found a function start in these bounds, something's wrong.
  assert (startfun >= 0), "ERROR in getFuncBounds::startfun"

  # Use a stack counter to find the closing brace:
  bracect = 0
  foundfirstbrace = False 
  endfun = startfun
  # ATTENTION: we've assumed ONLY ONE of each kind of brace per line!
  while (bracect > 0 or not foundfirstbrace):
    # If we've run off the end of the array, something's wrong.
    assert endfun < len(filelines) 
    if (filelines[endfun].find("{") >= 0):
      bracect = bracect + 1
      foundfirstbrace = True
    if (filelines[endfun].find("}") >= 0):
      bracect = bracect - 1
    endfun = endfun + 1

  funfile.close() 
  return Funcbounds(start=startfun, end=endfun) 


def stackwhack(cur_thread):
  """ Call thusly: ls.stackwhack(lldb.frame.GetThread())
      Args:
        cur_thread: The SBThread on which the stack to be whacked lives.

      Returns:
        nothing.
  """
  global stackfilename
  frame_ct = cur_thread.GetNumFrames()
  stack = []
  i = 0
  # Move up through the stack frames until we hit one without a line number.
  while (i < frame_ct):
    nextf = cur_thread.GetFrameAtIndex(i)
    lineentry = nextf.GetLineEntry()
    if (not lineentry):
      break
    funname = nextf.GetFunctionName()
    funtype = str(nextf.GetFunction().GetType())
    fullname = lineentry.GetFileSpec().GetDirectory() + "/" + \
               lineentry.GetFileSpec().GetFilename() 
    inclline = lineentry.GetLine() - 1 # adjust from 1- to 0-based.
    funbounds = getFuncBounds(funname, funtype, fullname, inclline) 
    # Push all information needed to print onto the stack.
    stack.append((fullname, inclline, funbounds.start, funbounds.end))
    i = i+1

  if (len(stack) < 1):
    print "No source code found in stack frames."
    return

  fileparts = cur_thread.GetFrameAtIndex(0).GetLineEntry().GetFileSpec().GetFilename().split(".")
  # Append filetype of innermost stack frame file, so syntax highlighting has a
  # fighting chance at getting things kind of right.
  fileext = fileparts[-1] if len(fileparts) > 1 else "txt"
  stackfilename = "/Users/mirandac/stackwhacker/" + stackfilename + "." + fileext
  stackfile = open(stackfilename, 'w')
  # TODO: print time at the top so we don't overwrite?

  # we need to change from while to recursive to do this properly...
  indent_level = 0
  while (len(stack) > 0): 
    fullname, inclline, startfun, endfun = stack.pop()
    funfile = open(fullname, 'r')
    filelines = funfile.read().splitlines()
    s = startfun 
    indent = " " * indent_level * 2
    # Write the name of the file as a comment starting each function.
    comment_indent = " " * ((indent_level * 2) - 2) # Negatives treated as 0.
    stackfile.write("//" + comment_indent + fullname + ":" + str(inclline) + "\n")
    while (s <= inclline):
      stackfile.write(indent + filelines[s])
      stackfile.write('\n')
      s = s+1
    stackfile.write('\n\n')
    funfile.close()
    indent_level = indent_level + 1
  stackfile.close()
  print "Wrote to " + stackfilename + "."


# See http://lldb.llvm.org/python-reference.html for details.
def whackme(debugger, command, result, internal_dict):
  """ To load, type command script import ~/stackwhacker/ls.py at the (lldb)
      prompt. Then the command "whackme" will be available at the prompt.
  """
  global stackfilename
  stackwhack(lldb.thread) 
  # TODO: if in Xcode, launch a new window with the file. How?
  result.PutCString('Stack whacked. See results in ./' + stackfilename + '.')


def __lldb_init_module(debugger, internal_dict):
  """ Installs 'whackme' as an lldb command.  To install:
        command script import <pathto>/stackwhack.py
  """
  debugger.HandleCommand('command script add -f ls.whackme whackme')
  print 'The "whackme" python command has been installed and is ready for use!'
